import numpy as np
import math
import pygad

# -------------------------- 1. 基础参数与复用函数 --------------------------
target_true = np.array([0, 200, 0])  # 真目标
r, h = 7, 10  # 真目标的半径7m、高10m
target_fake = np.array([0, 0, 0])  # 假目标
missile_M1 = np.array([20000, 0, 2000])  # M1初始位置
missile_speed = 300  # 导弹速度
drone_FY1 = np.array([17800, 0, 1800])  # FY1初始位置
g = 9.8  # 重力加速度
smoke_radius = 10  # 烟幕半径
smoke_sink_speed = 3  # 下沉速度
smoke_valid_time = 20  # 有效时间
drone_speed_range = [70, 140]  # 无人机速度约束
NUM_SMOKE = 3  # 烟幕弹数量

# 计算导弹飞到假目标的时间（关键约束）
missile_distance_to_fake = np.linalg.norm(target_fake - missile_M1)
missile_time_to_fake = missile_distance_to_fake / missile_speed  # 导弹到达假目标的时间


# 定义获取长方体顶点的函数
def get_cylinder_bounding_box_vertices(target_center, radius, height):
    x_min = target_center[0] - radius
    x_max = target_center[0] + radius
    y_min = target_center[1] - radius
    y_max = target_center[1] + radius
    z_min = target_center[2]
    z_max = target_center[2] + height
    #长方体顶点
    vertices = [
        np.array([x_min, y_min, z_min]),
        np.array([x_max, y_min, z_min]),
        np.array([x_max, y_max, z_min]),
        np.array([x_min, y_max, z_min]),
        np.array([x_min, y_min, z_max]),
        np.array([x_max, y_min, z_max]),
        np.array([x_max, y_max, z_max]),
        np.array([x_min, y_max, z_max])
    ]
    return vertices

# 预计算目标顶点
target_vertices = get_cylinder_bounding_box_vertices(target_true, r, h)

def missile_real_time_position(missile_init_pos, t):
    # 导弹到达假目标后不再移动
    if t >= missile_time_to_fake:
        return target_fake
    if t < 0:
        return missile_init_pos
    dir_vec = target_fake - missile_init_pos
    dir_norm = np.linalg.norm(dir_vec)
    if dir_norm <= 1e-6:
        return missile_init_pos
    dir_unit_vec = dir_vec / dir_norm
    return missile_init_pos + missile_speed * t * dir_unit_vec


# -------------------------- 2. 烟幕实时位置计算 --------------------------
def get_smoke_explosion_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate):
    # 投放点坐标
    drop_x = drone_init_pos[0] + drone_speed * math.cos(theta) * t_drop
    drop_y = drone_init_pos[1] + drone_speed * math.sin(theta) * t_drop
    drop_z = drone_init_pos[2]
    # 烟幕弹平抛运动
    explosion_x = drop_x + drone_speed * math.cos(theta) * t_detonate
    explosion_y = drop_y + drone_speed * math.sin(theta) * t_detonate
    # 竖直自由下落
    fall_distance = 0.5 * g * (t_detonate ** 2)
    explosion_z = max(drop_z - fall_distance, 0.0)
    return np.array([explosion_x, explosion_y, explosion_z])


def smoke_real_time_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate, t):
    detonate_t = t_drop + t_detonate
    if t < detonate_t - 1e-6:
        return np.array([0, 0, 0])
    explosion_center = get_smoke_explosion_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate)
    sink_distance = smoke_sink_speed * (t - detonate_t)
    real_z = max(explosion_center[2] - sink_distance, 0.0)
    return np.array([explosion_center[0], explosion_center[1], real_z])


# -------------------------- 3. 遮蔽判断 --------------------------
def is_segment_sphere_intersect(missile_pos, vertex_pos, smoke_center, smoke_radius):
    vec_AC = smoke_center - missile_pos
    vec_AB = vertex_pos - missile_pos
    ab_sq_mag = np.dot(vec_AB, vec_AB)
    if ab_sq_mag < 1e-12:
        return False
    k = np.dot(vec_AC, vec_AB) / ab_sq_mag
    if k <= 0.0:
        min_dist = np.linalg.norm(vec_AC)
    elif 0.0 < k < 1.0:
        cross_prod = np.cross(vec_AC, vec_AB)
        min_dist = np.linalg.norm(cross_prod) / np.sqrt(ab_sq_mag)
    else:
        vec_BC = smoke_center - vertex_pos
        min_dist = np.linalg.norm(vec_BC)
    return min_dist <= smoke_radius + 1e-6


def is_effective_occlusion(t, drone_speed, theta, t_drop, t_detonate):
    detonate_t = t_drop + t_detonate
    smoke_start = detonate_t
    smoke_end = detonate_t + smoke_valid_time
    if t < smoke_start - 1e-6 or t > smoke_end + 1e-6:
        return False
    missile_pos = missile_real_time_position(missile_M1, t)
    smoke_center = smoke_real_time_center(drone_FY1, drone_speed, theta, t_drop, t_detonate, t)
    for vertex in target_vertices:
        if is_segment_sphere_intersect(missile_pos, vertex, smoke_center, smoke_radius):
            return True
    return False


# 适应度函数：联合遮蔽时长
single_durations = [0.0 for _ in range(NUM_SMOKE)]
def fitness_func(ga_instance, solution, solution_idx):
    global single_durations
    params = []
    for i in range(NUM_SMOKE):
        base = i * 4
        drone_speed = solution[base]
        theta = solution[base+1]
        t_drop = solution[base+2]
        t_detonate = solution[base+3]
        params.append((drone_speed, theta, t_drop, t_detonate))
    # 统计每枚弹的单独干扰时长
    single_durations = []
    time_step = 0.01
    for i in range(NUM_SMOKE):
        drone_speed, theta, t_drop, t_detonate = params[i]
        detonate_t = t_drop + t_detonate
        smoke_start = detonate_t
        smoke_end = detonate_t + smoke_valid_time
        total_duration = 0.0
        prev_effective = False
        t = smoke_start
        while t <= smoke_end + 1e-6:
            current_effective = is_effective_occlusion(t, drone_speed, theta, t_drop, t_detonate)
            if current_effective and not prev_effective:
                start_t = t
            elif not current_effective and prev_effective:
                total_duration += t - start_t
            prev_effective = current_effective
            t += time_step
        if prev_effective:
            total_duration += smoke_end - start_t
        single_durations.append(round(total_duration, 3))
    # 统计联合遮蔽时长
    detonate_times = [p[2]+p[3] for p in params]
    smoke_start = min(detonate_times)
    smoke_end = max([dt+smoke_valid_time for dt in detonate_times])
    total_duration = 0.0
    prev_effective = False
    t = smoke_start
    while t <= smoke_end + 1e-6:
        current_effective = False
        for i in range(NUM_SMOKE):
            current_effective |= is_effective_occlusion(t, *params[i])
        if current_effective and not prev_effective:
            start_t = t
        elif not current_effective and prev_effective:
            total_duration += t - start_t
        prev_effective = current_effective
        t += time_step
    if prev_effective:
        total_duration += smoke_end - start_t
    return round(total_duration, 3)

# gene_space设置
gene_space = []
for i in range(NUM_SMOKE):
    gene_space += [
        {'low': 70, 'high': 140},  # 速度
        {'low': 0, 'high': 2 * math.pi},  # 方向角
        {'low': 0, 'high': 30},  # 投放延迟
        {'low': 0, 'high': 10}   # 起爆延迟
    ]

def on_generation(ga_instance):
    # 获取当前种群最大适应度
    best_fitness = np.max(ga_instance.last_generation_fitness)
    print(f"当前迭代次数：{ga_instance.generations_completed}, 最大遮蔽时长：{best_fitness:.2f}s")

def fix_t_drop_order(candidate, num_smoke):
    # 取出所有投放延迟
    t_drop_list = [candidate[i*4+2] for i in range(num_smoke)]
    # 排序并保证间隔
    t_drop_list_sorted = [t_drop_list[0]]
    for i in range(1, num_smoke):
        t_drop_next = max(t_drop_list[i], t_drop_list_sorted[-1] + 1.0)
        t_drop_list_sorted.append(t_drop_next)
    # 回填
    for i in range(num_smoke):
        candidate[i*4+2] = t_drop_list_sorted[i]
    return candidate

def on_crossover(ga_instance, offspring_crossover):
    for candidate in offspring_crossover:
        fix_t_drop_order(candidate, NUM_SMOKE)
    return offspring_crossover

def on_mutation(ga_instance, offspring_mutation):
    for candidate in offspring_mutation:
        fix_t_drop_order(candidate, NUM_SMOKE)
    return offspring_mutation

# 生成有效种群
def generate_valid_population(pop_size, num_genes, gene_space, num_smoke):
    population = []
    while len(population) < pop_size:
        candidate = []
        t_drop_list = []
        # 先生成第一个烟雾弹的投放延迟
        t_drop1 = np.random.uniform(gene_space[2]['low'], gene_space[2]['high'])
        t_drop_list.append(t_drop1)
        # 后续烟雾弹的投放延迟必须比前一个大至少1秒
        for i in range(1, num_smoke):
            min_t = t_drop_list[-1] + 1.0
            max_t = gene_space[i*4+2]['high']
            # 如果最大值比最小值还小，直接用最小值
            if max_t < min_t:
                t_drop = min_t
            else:
                t_drop = np.random.uniform(min_t, max_t)
            t_drop_list.append(t_drop)
        # 组装完整个体
        for i in range(num_smoke):
            candidate += [
                np.random.uniform(gene_space[i*4]['low'], gene_space[i*4]['high']),      # 速度
                np.random.uniform(gene_space[i*4+1]['low'], gene_space[i*4+1]['high']),  # 方向
                t_drop_list[i],                                                          # 投放延迟
                np.random.uniform(gene_space[i*4+3]['low'], gene_space[i*4+3]['high'])   # 起爆延迟
            ]
        population.append(candidate)
    return np.array(population)

# 遗传算法主程序
if __name__ == "__main__":
    initial_population = generate_valid_population(
        pop_size=50, num_genes=4*NUM_SMOKE, gene_space=gene_space, num_smoke=NUM_SMOKE
    )
    ga_instance = pygad.GA(
        num_generations=150,  # 迭代次数
        num_parents_mating=10,  # 每代参与交配的父代数量
        fitness_func=fitness_func,  # 适应度函数
        sol_per_pop=50, # 种群大小
        num_genes=4*NUM_SMOKE,  # 变量数量（12个优化参数）
        gene_space=gene_space, # 变量范围
        parent_selection_type="tournament",  # 选择策略：锦标赛选择
        crossover_type="two_points",  # 交叉策略：单点交叉
        mutation_type="random",  # 变异策略：随机变异
        mutation_percent_genes=10,  # 变异基因比例（10%）
        mutation_num_genes=4,  # 直接指定变异1个基因
        parallel_processing=64,  # 启用并行计算
        on_generation=on_generation,  # 每一代的回调函数
        on_crossover=on_crossover,  # 交叉操作完成后的回调函数
        on_mutation=on_mutation,    # 变异操作完成后的回调函数
        initial_population=initial_population, # 自定义初始种群
    )

    print("开始多枚烟幕弹联合优化...")
    ga_instance.run()
    solution, solution_fitness, solution_idx = ga_instance.best_solution()
    print("\n最优解：")
    for i in range(NUM_SMOKE):
        base = i * 4
        drone_speed = solution[base]
        theta = solution[base+1]
        t_drop = solution[base+2]
        t_detonate = solution[base+3]
        theta_deg = (theta * 180 / math.pi) % 360
        drop_x = drone_FY1[0] + drone_speed * math.cos(theta) * t_drop
        drop_y = drone_FY1[1] + drone_speed * math.sin(theta) * t_drop
        drop_z = drone_FY1[2]
        explosion_point = get_smoke_explosion_center(
            drone_FY1, drone_speed, theta, t_drop, t_detonate
        )
        print(f"烟幕弹{i+1}: 速度={drone_speed:.2f}m/s, 方向角={theta_deg:.2f}°, 投放时间：{t_drop:.2f}s , 起爆延迟={t_detonate:.2f}s")
        print(f"  投放点=({drop_x:.2f}, {drop_y:.2f}, {drop_z:.2f})")
        print(f"  起爆点=({explosion_point[0]:.2f}, {explosion_point[1]:.2f}, {explosion_point[2]:.2f})")
        print(f"  有效干扰时长：{single_durations[i]:.2f}s")
    print(f"联合最大遮蔽时长：{solution_fitness:.2f}s")

